import os
import xml.etree.ElementTree as ET


def convert(xml_files, txt_dir, class_list):
    for xml in xml_files:
        tree = ET.parse(xml)
        root = tree.getroot()
        filename = root.find("filename").text.split('.')[0]
        width = int(root.find("size").find("width").text)
        height = int(root.find("size").find("height").text)
        box_list = []
        for obj in root.findall("object"):
            name = obj.find("name").text
            class_id = class_list.index(name)
            bondbox = obj.find("bndbox")
            x_min = int(bondbox.find("xmin").text)
            x_max = int(bondbox.find("xmax").text)
            y_min = int(bondbox.find("ymin").text)
            y_max = int(bondbox.find("ymax").text)

            x_center = (x_min + x_max) / (2 * width)
            y_center = (y_min + y_max) / (2 * height)
            box_width = (x_max - x_min) / width
            box_height = (y_max - y_min) / height

            line = str(class_id) + " " + str(x_center) + " " + str(y_center) + " " + str(box_width) + " " + str(box_height)
            box_list.append(line)
        with open(os.path.join(txt_dir, filename + '.txt'), 'w', encoding='utf-8') as file:
            for line in box_list:
                file.write(line)
                if box_list.index(line) != len(box_list) - 1:
                    file.write("\n")


if __name__ == '__main__':
    xml_dir = "D:/24926/Desktop/MyProject/thesis/数据集/dataset/labels/val"
    txt_dir = "D:/24926/Desktop/MyProject/thesis/数据集/dataset/labels/val"
    class_list = ['Crack', 'Dent', 'Rust']
    items = os.listdir(xml_dir)
    xml_files = []
    for item in items:
        if item.endswith('.xml'):
            xml_files.append(os.path.join(xml_dir, item))
    convert(xml_files, txt_dir, class_list)
